# -*- coding: utf-8 -*-
"""
@Time    : 2024/6/27 10:52 
@Author  : ZhangShenao 
@File    : 2.JsonOutputParser组件使用.py
@Desc    : JsonOutputParser组件使用
"""
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.output_parsers import JsonOutputParser
from pydantic.v1 import BaseModel, Field


# 构造一个json数据结构,放在Prompt指令中,用于告诉LLM输出的格式
class Joke(BaseModel):
    """
    Joke 数据结构,放在Prompt指令中,用于告诉LLM输出的格式
    """

    joke: str = Field(description='the content of the joke')
    punchline: str = Field(description='the punchline of the joke')


# 创建智普AI LLM
llm = ChatZhipuAI(
    model="GLM-4-Air",
    temperature=0.5,
)

# 创建JsonOutputParser输出解析器,并指定输出格式
parser = JsonOutputParser(pydantic_object=Joke)
print(parser.get_format_instructions())

# 编排PromptTemplate
# chat_prompt_template = ChatPromptTemplate.from_messages(
#     [
#         ('system', 'You are a chat bot developed by ZhiPu AI,please answer user\'s questions.'),
#         ('human', '{instrument}\n'
#                   'user\'s question：{query}')
#     ]
# ).partial(instrument=parser.get_format_instructions())
#
# # 生成Prompt,调用LLM获取响应,并解析输出结果
# content = parser.invoke(llm.invoke(chat_prompt_template.invoke(
#     {'query': 'Please tell me a joke about programmers'}
# )))
#
# print(type(content))
# print(content.get('joke'))
# print(content.get('punchline'))
